'''
Author: jojo
Date: 2021-06-15 06:40:34
LastEditors: jojo
LastEditTime: 2021-06-18 03:07:28
FilePath: /waifu/utils/video_utils.py
'''
import subprocess
import os
from tqdm import tqdm
import shutil
# from interpolation import get_fps,read_video_info
import signal
from . import global_var_manager as gv
import sys
import json


preset = 9

def getVideoTime(path):
    cmdline = 'ffprobe "%s" -show_entries format=duration -of compact=p=0:nk=1 -v 0'%path
    gettime=subprocess.check_output(cmdline, shell=True)
    timeT=int(float(gettime.strip()))
    return timeT

def read_video_info(src_dir,filename):
    res_json = os.path.join(src_dir,filename + ".json")
    cmd = f'ffprobe -v quiet -print_format json -show_format -show_streams {os.path.join(src_dir,filename)} >> {res_json}'
    # command = 'ffprobe -v error -select_streams v -show_entries stream=nb_frames -of json {os.path.join(src_dir,filename)} >> {res_json}'
    os.system(cmd)
    res_dict  = {}
    with open(res_json,'r') as file:
        res_dict = json.load(file)  
    
    os.remove(res_json)
    return res_dict

def get_fps(res_dict):
    content = res_dict['streams'][0]['avg_frame_rate']
    up,down = content.split('/')
    fps = float(up) / float(down)
    return fps

def cut(video_path,cut_time=600):
    """cutting the video

    Args:
        video_path (str): the video path
        cut_time (int, optional): clip length. Defaults to 600.

    Returns:
        clip_path_list(list): clip list(sorted)
        output_abs_root_path(list): output root absolute path
        clip_name_list(list): clip name list(inlcuding suffix)
    """
    signal.signal(signal.SIGINT,interupt_handler)
    signal.signal(signal.SIGHUP, interupt_handler)
    signal.signal(signal.SIGTERM, interupt_handler)
    # 验证文件是否存在
    if not os.path.exists(video_path):
        print(f'{video_path} does not exist!')
        return None 
    total_time=getVideoTime(video_path)
    start_time=0
    index=1
    filename_fullname = video_path.split('/')[-1]
    filename,suffix = filename_fullname.split('.')
    output_path = os.path.join(video_path.split(filename)[0],filename+'-cut')
    if not os.path.exists(output_path):
        os.mkdir(output_path)
    clip_path_list = []
    clip_name_list = []
    gv.set_value('clip_out',output_path)
    with tqdm(total=total_time,desc='cutting video') as pbar:
        while start_time<total_time:
            # cmdLine = 'ffmpeg -ss %s -i %s -c copy -t %s %s.mp4 -loglevel quiet -y'%(firstTime,video_path,cut_time,'cut_%s'%index)
            clip_name = f'{filename}-{index}.{suffix}'
            current_path = f'{output_path}/{clip_name}'

            # 直接分割
            cmdLine = f'ffmpeg -threads 10 -ss {start_time} -i {video_path} -c copy -t {cut_time} {current_path} -loglevel quiet -y'

            # CPU 转码
            # cmdLine = f'ffmpeg -ss {start_time} -t {cut_time} -i {video_path} -c:v libx264  {current_path} -loglevel quiet -y'

            # GPU 转码
            # 解决方案：https://www.jianshu.com/p/00b6de770b1a?utm_campaign=maleskine
            # cmdLine = f'ffmpeg -ss {start_time} -t {cut_time} -c:v h264_cuvid -i {video_path} -c:v h264_nvenc -c:a aac  {current_path} -loglevel quiet -y' # fixme 转码之后才能concat，但事实上有没有必要？


            clip_path_list.append(os.path.abspath(current_path))
            clip_name_list.append(clip_name)
            returnCmd = subprocess.call(cmdLine, shell=True)
            start_time+=cut_time
            index+=1
            pbar.update(cut_time)
            
    output_abs_root_path = os.path.abspath(output_path)
    
    return clip_path_list,output_abs_root_path,clip_name_list

def concat(video_path_list,video_full_name,output_path):
    """concat the interpolated clip to destination path

    Args:
        video_path_list (list): clip path list
        video_full_name (str): original video full name(including suffix)
        output_path (str): destination path

    Returns:
        int: cmd result
    """
    if len(video_path_list)<1:
        print("Empty video path list")
        return None
    
    # video_full_name = src_path.split('/')[-1]
    video_name,suffix = video_full_name.split('.')
    # 写入文件中
    filename = video_path_list[0].split('/')[-1]
    file_root = video_path_list[0].split(filename)[0]
    file_list = os.path.join(file_root,video_name+'.txt')
    gv.set_value('file_list',file_list)
    with open(file_list,'w') as file:
        for video_path in video_path_list:
            text = f"file '{video_path}'\n"
            file.write(text)
    
    cmd = f'ffmpeg -f concat -safe 0 -i {file_list} -c copy {output_path}/{video_name}.mp4 -loglevel quiet -y'
    print('concating video...')
    
    returnCmd = subprocess.call(cmd, shell=True)
    
    remove_file(file_list)
    
    return returnCmd

def extract_audio(src_dir,filename):
    file_realname,_ = filename.split('.')
    src_file_path = os.path.join(src_dir,filename)
    audio_path = os.path.join(src_dir,file_realname + '.m4a')
    # extract audio
    ext_aud_com = f'ffmpeg -y -i {src_file_path} -vn -acodec copy {audio_path}'
    os.system(ext_aud_com)
    
    return audio_path

def cont_audio(src_path,filename,audio_path,dst_file_path):
    
    res_dict = read_video_info(src_dir=src_path,filename=filename)
    fps = get_fps(res_dict)
    pix_fmt = res_dict['streams'][0]['pix_fmt']
    
    video_path = os.path.join(src_path,filename)
    
    # how to determine the crf value: https://zhuanlan.zhihu.com/p/250590703
    # int_with_aud_com = f'ffmpeg -y -framerate {double_fps} -i {output_frame_dir}/%08d.png -i {audio_path} -c:a copy -crf 23 -c:v libx264 -pix_fmt {pix_fmt} {dst_file_path}'

    # int_with_aud_com = f'ffmpeg -y -framerate {double_fps} -i {video_path} -i {audio_path} -c:a copy -crf 23 -c:v h264_nvenc -pix_fmt {pix_fmt} {dst_file_path}' # use gpu

    # use gpu but seems not so good
    # int_with_aud_com = f'ffmpeg -y -i {video_path} -i {audio_path} -c:a copy -c:v h264_nvenc -preset {preset} -pix_fmt {pix_fmt} -r {fps} {dst_file_path}/{filename}' # use gpu

    # use cpu
    int_with_aud_com = f'ffmpeg -y -i {video_path} -i {audio_path} -c:a copy -crf 23 -c:v libx264 -pix_fmt {pix_fmt} -r {fps} {dst_file_path}/{filename}' 
    os.system(int_with_aud_com)
    
def remove_path(dir):
    try:
        shutil.rmtree(dir)
    except Exception as e:
        print(f"{dir} does not exist!")
        
def remove_file(filepath):
    try:
        os.remove(filepath)
    except Exception as e:
        print(f"{filepath} does not exist!")
        
def interupt_handler(signum,frame):
    print("user has interupt the progress")
    clip_output_path = gv.get_value('clip_out')
    file_list = gv.get_value('file_list')
    print("cleaning cache...")
    remove_path(clip_output_path)
    remove_file(file_list)
    sys.exit()

        
if __name__=='__main__':
    video_path = '../workspace/src/p2/test.mp4'
    clip_list,_,_ = cut(video_path=video_path)
    output_path = '../workspace/dst/test'
    concat(clip_list,video_path,output_path)